# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/16 14:41 
@Author  : ZhangShenao 
@File    : 4.使用ConversationSummaryMemory对话摘要记忆.py
@Desc    : 使用ConversationSummaryMemory,实现对话摘要记忆

ConversationSummaryMemory会将之前的对话历史
调用LLM的接口生成摘要总结
然后通过{history}传递到当前Prompt中

对于长对话,ConversationSummaryMemory可以有效生成长期记忆,并且节省每轮对话消耗的Token数量
"""

import os
from operator import itemgetter

import dotenv
from langchain.memory import ConversationSummaryMemory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI

# 系统指令
SYSTEM_PROMPT = '你是一个具有记忆功能的AI助手，请根据结合用户的当前提问，和之前的历史聊天记录，生成准确的回答。你们的对话全部采用中文'

# 聊天历史Key
HISTORY_KEY = 'history'

# 加载环境变量
dotenv.load_dotenv()

# 构造LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))

# 创建ConversationSummaryMemory,对话总结记忆组件
memory = ConversationSummaryMemory(
    llm=llm,  # 指定用于生成总结的LLM
    memory_key=HISTORY_KEY,  # 指定聊天历史Key
    input_key='input',  # 指定人类输入Key
    output_key='output',  # 指定模型输出Key
    return_messages=True,  # 结果返回聊天消息列表,而不是字符串
)

# 构造Prompt
prompt = ChatPromptTemplate.from_messages([
    ('system', SYSTEM_PROMPT),
    MessagesPlaceholder(HISTORY_KEY),  # 使用MessagesPlaceholder消息占位符,占位{history}聊天历史
    ('human', '{input}')
])

# 构造OutputParser
parser = StrOutputParser()

# 构建Chain
# 使用RunnablePassthrough.assign()函数,为history参数赋值
# 使用RunnableLambda将一个lambda表达式封装成Runnable可运行协议
# 使用itemgetter()函数,从memory返回的dict中获取history的值,返回一个字符串
chain = RunnablePassthrough.assign(
    history=RunnableLambda(memory.load_memory_variables) | itemgetter(HISTORY_KEY)
) | prompt | llm | parser

# 接收用户输入
while True:
    human_input = input('Human：')
    if 'bye' == human_input:
        print('bye bye~')
        break

    # 按照stream模式,执行chain,并打印输出
    output_chunks = chain.stream({'input': human_input})
    output = ''
    print('AI：', end='', flush=True)
    for chunk in output_chunks:
        output += chunk
        print(chunk, end='', flush=True)
    print()

    # 保存当前轮的聊天记录
    memory.save_context(inputs={'input': human_input}, outputs={'output': output})

    # 打印历史聊天记录
    print(f'聊天历史: {memory.load_memory_variables({})}')
